{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.vision import *\n",
    "import warnings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def res_block(nf, depth=1, dense:bool=False, norm_type:Optional[NormType]=NormType.Batch, **conv_kwargs):\n",
    "    \"Resnet block of `nf` features. `conv_kwargs` are passed to `conv_layer`.\"\n",
    "    norm2 = norm_type\n",
    "    if not dense and (norm_type==NormType.Batch): norm2 = NormType.BatchZero\n",
    "    layers = []\n",
    "    for i in range(depth+1):\n",
    "        layers.append(conv_layer(nf, nf, norm_type=norm2, **conv_kwargs))\n",
    "    return SequentialEx(*layers)\n",
    "\n",
    "def res_block_with_depth(n_in, n_out, n_basic_blocks=1, depth=1, dense:bool=False, norm_type:Optional[NormType]=NormType.Batch, **conv_kwargs):\n",
    "    \"Resnet block of `nf` features. `conv_kwargs` are passed to `conv_layer`.\"\n",
    "    norm2 = norm_type\n",
    "    if not dense and (norm_type==NormType.Batch): norm2 = NormType.BatchZero\n",
    "    layer1 = [conv_layer(n_in, n_out, norm_type=norm_type, stride=(2, 2), **conv_kwargs)]\n",
    "    if n_basic_blocks == 0 : return layer1[0]\n",
    "    layers = [res_block(n_in, depth=depth) for _ in range(n_basic_blocks)]\n",
    "    return SequentialEx(*(layers+layer1))\n",
    "\n",
    "def get_resnet(input_channels=3, out_classes=10, n_basic_blocks=1 ,depth=2, n_blocks=3, fc=True):\n",
    "    \"\"\"\n",
    "    :param input_channels: Number of input channels (i.e Image channels)\n",
    "    :param out_classes: Number of output classes\n",
    "    :param n_basic_blocks: Number of basic blocks after one downsampling conv layer (child block/basic block)\n",
    "    :param depth: depth of each child block / basic block\n",
    "    :param n_block: Number of blocks with downsampling conv layers (parent block)\n",
    "    :return: resnet\n",
    "    \"\"\"\n",
    "    if depth > 2: warnings.warn(\"Do not use this function for more than 2 depth, The residual connections will become insufficient\")\n",
    "    net = [conv_layer(3, 64, ks = 7, stride = 2, padding = 3),\n",
    "            nn.MaxPool2d(3, 2, padding = 1),\n",
    "            res_block(64, depth=depth)]\n",
    "    in_c = net[0][0].out_channels\n",
    "    for i in range(n_blocks):\n",
    "        net.append(res_block_with_depth(in_c, 2*in_c,n_basic_blocks=n_basic_blocks, depth=depth))\n",
    "        in_c *= 2\n",
    "    if fc: net.append(create_head(in_c*2, 10, ps=0))\n",
    "    return SequentialEx(*net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Incorrect implementation : DON'T USE\n",
    "def bottleneck(prev_c, nf, downsample:bool = False, downsample_first_block:bool = False, dense:bool=False, norm_type:Optional[NormType]=NormType.Batch, **conv_kwargs) :\n",
    "    norm2 = norm_type\n",
    "    if not dense and (norm_type==NormType.Batch): norm2 = NormType.BatchZero\n",
    "    layers = []\n",
    "    layers.append(conv_layer(prev_c, nf, 1, norm_type = norm2, **conv_kwargs))\n",
    "    layers.append(conv_layer(nf, nf, 3, norm_type = norm2, **conv_kwargs))\n",
    "    layers.append(conv_layer(nf, nf * 4, 1, norm_type = norm2, **conv_kwargs))\n",
    "    if downsample : \n",
    "        layers.append(conv_layer(prev_c, nf * 4, 1, stride = 2, norm_type = norm2, **conv_kwargs))\n",
    "    elif downsample_first_block : \n",
    "        layers.append(conv_layer(prev_c, nf * 4, 1, stride = 1, norm_type = norm2, **conv_kwargs))\n",
    "    layers.append(MergeLayer(dense))\n",
    "    return SequentialEx(*layers)\n",
    "\n",
    "def get_resnet50(in_c = 3, n_out = 10, n_blocks = 4, n_basic_blocks = 3, dense:bool = False, fc:bool = True, norm_type:Optional[NormType] = NormType.Batch, **conv_kwargs) : \n",
    "    net = [conv_layer(3, 64, ks = 7, stride = 2, padding = 3),\n",
    "           nn.MaxPool2d(3, 2, padding = 1)\n",
    "          ]\n",
    "    count = 0\n",
    "    in_c = 64\n",
    "    prev_c = net[0][0].out_channels\n",
    "    for i in range(n_blocks) : \n",
    "        for j in range(n_basic_blocks) : \n",
    "            if j == 0 and i == 0 :\n",
    "                net.append(bottleneck(prev_c, in_c, downsample_first_block = True))\n",
    "            elif j == 0 : \n",
    "                net.append(bottleneck(prev_c, in_c, downsample = True))\n",
    "            else : \n",
    "                net.append(bottleneck(prev_c, in_c, downsample = False))\n",
    "            prev_c = net[count + 2][-2][-1].num_features\n",
    "            print(prev_c)\n",
    "            count += 1\n",
    "        in_c *= 2\n",
    "    \n",
    "    if fc : \n",
    "        net.append(create_head(in_c * 2, 10, ps = 0))\n",
    "    return SequentialEx(*net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequentialEx(\n",
       "  (layers): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "      (1): ReLU(inplace)\n",
       "      (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "    (2): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (3): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (1): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (2): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (3): Sequential(\n",
       "          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (4): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (1): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (2): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (3): Sequential(\n",
       "          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (5): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (1): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (2): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (3): Sequential(\n",
       "          (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (6): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (1): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (2): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (3): Sequential(\n",
       "          (0): Conv2d(512, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (7): Sequential(\n",
       "      (0): AdaptiveConcatPool2d(\n",
       "        (ap): AdaptiveAvgPool2d(output_size=1)\n",
       "        (mp): AdaptiveMaxPool2d(output_size=1)\n",
       "      )\n",
       "      (1): Flatten()\n",
       "      (2): BatchNorm1d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (3): Linear(in_features=2048, out_features=512, bias=True)\n",
       "      (4): ReLU(inplace)\n",
       "      (5): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (6): Linear(in_features=512, out_features=10, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get_resnet(n_blocks = 4, n_basic_blocks = 3).cuda().forward(torch.randn(2, 3, 224, 224).cuda()).shape\n",
    "get_resnet(n_blocks = 4, n_basic_blocks = 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 112, 112]           9,408\n",
      "       BatchNorm2d-2         [-1, 64, 112, 112]             128\n",
      "              ReLU-3         [-1, 64, 112, 112]               0\n",
      "         MaxPool2d-4           [-1, 64, 56, 56]               0\n",
      "            Conv2d-5           [-1, 64, 56, 56]           4,096\n",
      "       BatchNorm2d-6           [-1, 64, 56, 56]             128\n",
      "              ReLU-7           [-1, 64, 56, 56]               0\n",
      "            Conv2d-8           [-1, 64, 56, 56]          36,864\n",
      "       BatchNorm2d-9           [-1, 64, 56, 56]             128\n",
      "             ReLU-10           [-1, 64, 56, 56]               0\n",
      "           Conv2d-11          [-1, 256, 56, 56]          16,384\n",
      "      BatchNorm2d-12          [-1, 256, 56, 56]             512\n",
      "           Conv2d-13          [-1, 256, 56, 56]          16,384\n",
      "      BatchNorm2d-14          [-1, 256, 56, 56]             512\n",
      "             ReLU-15          [-1, 256, 56, 56]               0\n",
      "       Bottleneck-16          [-1, 256, 56, 56]               0\n",
      "           Conv2d-17           [-1, 64, 56, 56]          16,384\n",
      "      BatchNorm2d-18           [-1, 64, 56, 56]             128\n",
      "             ReLU-19           [-1, 64, 56, 56]               0\n",
      "           Conv2d-20           [-1, 64, 56, 56]          36,864\n",
      "      BatchNorm2d-21           [-1, 64, 56, 56]             128\n",
      "             ReLU-22           [-1, 64, 56, 56]               0\n",
      "           Conv2d-23          [-1, 256, 56, 56]          16,384\n",
      "      BatchNorm2d-24          [-1, 256, 56, 56]             512\n",
      "             ReLU-25          [-1, 256, 56, 56]               0\n",
      "       Bottleneck-26          [-1, 256, 56, 56]               0\n",
      "           Conv2d-27          [-1, 128, 56, 56]          32,768\n",
      "      BatchNorm2d-28          [-1, 128, 56, 56]             256\n",
      "             ReLU-29          [-1, 128, 56, 56]               0\n",
      "           Conv2d-30          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-31          [-1, 128, 28, 28]             256\n",
      "             ReLU-32          [-1, 128, 28, 28]               0\n",
      "           Conv2d-33          [-1, 512, 28, 28]          65,536\n",
      "      BatchNorm2d-34          [-1, 512, 28, 28]           1,024\n",
      "           Conv2d-35          [-1, 512, 28, 28]         131,072\n",
      "      BatchNorm2d-36          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-37          [-1, 512, 28, 28]               0\n",
      "       Bottleneck-38          [-1, 512, 28, 28]               0\n",
      "           Conv2d-39          [-1, 128, 28, 28]          65,536\n",
      "      BatchNorm2d-40          [-1, 128, 28, 28]             256\n",
      "             ReLU-41          [-1, 128, 28, 28]               0\n",
      "           Conv2d-42          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-43          [-1, 128, 28, 28]             256\n",
      "             ReLU-44          [-1, 128, 28, 28]               0\n",
      "           Conv2d-45          [-1, 512, 28, 28]          65,536\n",
      "      BatchNorm2d-46          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-47          [-1, 512, 28, 28]               0\n",
      "       Bottleneck-48          [-1, 512, 28, 28]               0\n",
      "           Conv2d-49          [-1, 256, 28, 28]         131,072\n",
      "      BatchNorm2d-50          [-1, 256, 28, 28]             512\n",
      "             ReLU-51          [-1, 256, 28, 28]               0\n",
      "           Conv2d-52          [-1, 256, 14, 14]         589,824\n",
      "      BatchNorm2d-53          [-1, 256, 14, 14]             512\n",
      "             ReLU-54          [-1, 256, 14, 14]               0\n",
      "           Conv2d-55         [-1, 1024, 14, 14]         262,144\n",
      "      BatchNorm2d-56         [-1, 1024, 14, 14]           2,048\n",
      "           Conv2d-57         [-1, 1024, 14, 14]         524,288\n",
      "      BatchNorm2d-58         [-1, 1024, 14, 14]           2,048\n",
      "             ReLU-59         [-1, 1024, 14, 14]               0\n",
      "       Bottleneck-60         [-1, 1024, 14, 14]               0\n",
      "           Conv2d-61          [-1, 256, 14, 14]         262,144\n",
      "      BatchNorm2d-62          [-1, 256, 14, 14]             512\n",
      "             ReLU-63          [-1, 256, 14, 14]               0\n",
      "           Conv2d-64          [-1, 256, 14, 14]         589,824\n",
      "      BatchNorm2d-65          [-1, 256, 14, 14]             512\n",
      "             ReLU-66          [-1, 256, 14, 14]               0\n",
      "           Conv2d-67         [-1, 1024, 14, 14]         262,144\n",
      "      BatchNorm2d-68         [-1, 1024, 14, 14]           2,048\n",
      "             ReLU-69         [-1, 1024, 14, 14]               0\n",
      "       Bottleneck-70         [-1, 1024, 14, 14]               0\n",
      "           Conv2d-71          [-1, 512, 14, 14]         524,288\n",
      "      BatchNorm2d-72          [-1, 512, 14, 14]           1,024\n",
      "             ReLU-73          [-1, 512, 14, 14]               0\n",
      "           Conv2d-74            [-1, 512, 7, 7]       2,359,296\n",
      "      BatchNorm2d-75            [-1, 512, 7, 7]           1,024\n",
      "             ReLU-76            [-1, 512, 7, 7]               0\n",
      "           Conv2d-77           [-1, 2048, 7, 7]       1,048,576\n",
      "      BatchNorm2d-78           [-1, 2048, 7, 7]           4,096\n",
      "           Conv2d-79           [-1, 2048, 7, 7]       2,097,152\n",
      "      BatchNorm2d-80           [-1, 2048, 7, 7]           4,096\n",
      "             ReLU-81           [-1, 2048, 7, 7]               0\n",
      "       Bottleneck-82           [-1, 2048, 7, 7]               0\n",
      "AdaptiveMaxPool2d-83           [-1, 2048, 1, 1]               0\n",
      "AdaptiveAvgPool2d-84           [-1, 2048, 1, 1]               0\n",
      "AdaptiveConcatPool2d-85           [-1, 4096, 1, 1]               0\n",
      "           Linear-86                 [-1, 1024]       4,195,328\n",
      "           Linear-87                   [-1, 10]          10,250\n",
      "================================================================\n",
      "Total params: 13,689,162\n",
      "Trainable params: 13,689,162\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.57\n",
      "Forward/backward pass size (MB): 177.12\n",
      "Params size (MB): 52.22\n",
      "Estimated Total Size (MB): 229.92\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "from models.custom_resnet import _resnet, Bottleneck\n",
    "net = _resnet('resnet50', Bottleneck, [2, 2, 2, 1], pretrained = False, progress = False).cuda()\n",
    "# print(net)\n",
    "from torchsummary import summary\n",
    "summary(net, (3, 224, 224))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models.resnets import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "2\n",
      "2\n",
      "2\n",
      "2\n",
      "2\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "SequentialEx(\n",
       "  (layers): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "      (1): ReLU(inplace)\n",
       "      (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "    (2): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): MergeLayer()\n",
       "      )\n",
       "    )\n",
       "    (3): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (4): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (5): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (6): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(512, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (7): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(1024, 2048, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (8): Sequential(\n",
       "      (0): AdaptiveConcatPool2d(\n",
       "        (ap): AdaptiveAvgPool2d(output_size=1)\n",
       "        (mp): AdaptiveMaxPool2d(output_size=1)\n",
       "      )\n",
       "      (1): Flatten()\n",
       "      (2): BatchNorm1d(4096, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (3): Linear(in_features=4096, out_features=512, bias=True)\n",
       "      (4): ReLU(inplace)\n",
       "      (5): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (6): Linear(in_features=512, out_features=10, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_resnet(n_blocks=5, n_basic_blocks=1, depth=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequentialEx(\n",
       "  (layers): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(1, 2, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (1): ReLU(inplace)\n",
       "      (2): BatchNorm2d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(2, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(2, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): MergeLayer()\n",
       "      )\n",
       "    )\n",
       "    (2): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(2, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(2, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): MergeLayer()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_block_with_depth(1, 2, n_blocks=2, depth=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequentialEx(\n",
       "  (layers): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(1, 1, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (1): ReLU(inplace)\n",
       "      (2): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(1, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): MergeLayer()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_block_with_depth(1, 1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_net = nn.Sequential(\n",
    "    conv_layer(3, 64, ks = 7, stride = 2, padding = 3),\n",
    "    nn.MaxPool2d(3, 2, padding = 1),\n",
    "    res_block(64, depth=2),\n",
    "    res_block_with_depth(64, 128, 2),\n",
    "    res_block_with_depth(128, 265, 2),\n",
    "    res_block_with_depth(256, 512, 2),\n",
    "    create_head(1024, 10, ps=0)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_resnet(input_channels=3, out_classes=10, depth=2, n_blocks=3):\n",
    "    net = [conv_layer(3, 64, ks = 7, stride = 2, padding = 3),\n",
    "            nn.MaxPool2d(3, 2, padding = 1),\n",
    "            res_block(64, depth=2)]\n",
    "    in_c = net[0][0].out_channels\n",
    "    for i in range(n_blocks):\n",
    "        net.append(res_block_with_depth(in_c, 2*in_c, depth=depth))\n",
    "        in_c *= 2\n",
    "    net.append(create_head(in_c*2, 10, ps=0))\n",
    "    return SequentialEx(*net)\n",
    "\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequentialEx(\n",
       "  (layers): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "      (1): ReLU(inplace)\n",
       "      (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "    (2): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): MergeLayer()\n",
       "      )\n",
       "    )\n",
       "    (3): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (4): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (5): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (6): Sequential(\n",
       "      (0): AdaptiveConcatPool2d(\n",
       "        (ap): AdaptiveAvgPool2d(output_size=1)\n",
       "        (mp): AdaptiveMaxPool2d(output_size=1)\n",
       "      )\n",
       "      (1): Flatten()\n",
       "      (2): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (3): Linear(in_features=1024, out_features=512, bias=True)\n",
       "      (4): ReLU(inplace)\n",
       "      (5): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (6): Linear(in_features=512, out_features=10, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_resnet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "64"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "la.out_channels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 10])"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nn.Sequential(*net).cuda().forward(torch.randn(2, 3, 224, 224).cuda()).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequentialEx(\n",
       "  (layers): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(265, 265, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (1): ReLU(inplace)\n",
       "      (2): BatchNorm2d(265, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): Sequential(\n",
       "      (0): Conv2d(265, 265, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (1): ReLU(inplace)\n",
       "      (2): BatchNorm2d(265, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (2): MergeLayer()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "    res_block(265)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequentialEx(\n",
       "  (layers): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(14, 7, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (1): ReLU(inplace)\n",
       "      (2): BatchNorm2d(7, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): Sequential(\n",
       "      (0): Conv2d(7, 14, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (1): ReLU(inplace)\n",
       "      (2): BatchNorm2d(14, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (2): MergeLayer()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_block(14, bottle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def res_block(n_in, n_out, dense:bool=False, norm_type:Optional[NormType]=NormType.Batch, **conv_kwargs):\n",
    "    \"Resnet block of `nf` features. `conv_kwargs` are passed to `conv_layer`.\"\n",
    "    norm2 = norm_type\n",
    "    if not dense and (norm_type==NormType.Batch): norm2 = NormType.BatchZero\n",
    "    return SequentialEx(conv_layer(n_in, n_out, norm_type=norm_type, **conv_kwargs),\n",
    "                      conv_layer(n_out, nf, norm_type=norm2, **conv_kwargs),\n",
    "                      MergeLayer(dense))\n",
    "\n",
    "res_block(n_in=64, n_out=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): AdaptiveConcatPool2d(\n",
       "    (ap): AdaptiveAvgPool2d(output_size=1)\n",
       "    (mp): AdaptiveMaxPool2d(output_size=1)\n",
       "  )\n",
       "  (1): Flatten()\n",
       "  (2): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (3): Dropout(p=0.25)\n",
       "  (4): Linear(in_features=1024, out_features=512, bias=True)\n",
       "  (5): ReLU(inplace)\n",
       "  (6): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (7): Dropout(p=0.5)\n",
       "  (8): Linear(in_features=512, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "ls = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(ls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.randn(3, 224, 224)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "adaptive_pooling = nn.AdaptiveMaxPool2d(output_size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[4.3917]],\n",
       "\n",
       "        [[3.8662]],\n",
       "\n",
       "        [[4.1667]]])"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adaptive_pooling(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(3.8662)"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a[1].max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models.resnets import get_resnet\n",
    "from models.unet import DynamicUnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = get_resnet(fc=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequentialEx(\n",
       "  (layers): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "      (1): ReLU(inplace)\n",
       "      (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "    (2): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): MergeLayer()\n",
       "      )\n",
       "    )\n",
       "    (3): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (4): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (5): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "unet = DynamicUnet(net, n_classes=10).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.vision import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DynamicUnet(\n",
       "  (layers): ModuleList(\n",
       "    (0): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "        (2): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): Sequential(\n",
       "              (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (2): MergeLayer()\n",
       "          )\n",
       "        )\n",
       "        (3): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): SequentialEx(\n",
       "              (layers): ModuleList(\n",
       "                (0): Sequential(\n",
       "                  (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "                  (1): ReLU(inplace)\n",
       "                  (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "                )\n",
       "                (1): MergeLayer()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (4): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): SequentialEx(\n",
       "              (layers): ModuleList(\n",
       "                (0): Sequential(\n",
       "                  (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "                  (1): ReLU(inplace)\n",
       "                  (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "                )\n",
       "                (1): MergeLayer()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (5): SequentialEx(\n",
       "          (layers): ModuleList(\n",
       "            (0): Sequential(\n",
       "              (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "              (1): ReLU(inplace)\n",
       "              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            )\n",
       "            (1): SequentialEx(\n",
       "              (layers): ModuleList(\n",
       "                (0): Sequential(\n",
       "                  (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "                  (1): ReLU(inplace)\n",
       "                  (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "                )\n",
       "                (1): MergeLayer()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU()\n",
       "    (3): Sequential(\n",
       "      (0): Sequential(\n",
       "        (0): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (1): Sequential(\n",
       "        (0): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (4): UnetBlock(\n",
       "      (shuf): PixelShuffle_ICNR(\n",
       "        (conv): Sequential(\n",
       "          (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (shuf): PixelShuffle(upscale_factor=2)\n",
       "        (pad): ReplicationPad2d((1, 0, 1, 0))\n",
       "        (blur): AvgPool2d(kernel_size=2, stride=1, padding=0)\n",
       "        (relu): ReLU(inplace)\n",
       "      )\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (relu): ReLU()\n",
       "    )\n",
       "    (5): UnetBlock(\n",
       "      (shuf): PixelShuffle_ICNR(\n",
       "        (conv): Sequential(\n",
       "          (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (shuf): PixelShuffle(upscale_factor=2)\n",
       "        (pad): ReplicationPad2d((1, 0, 1, 0))\n",
       "        (blur): AvgPool2d(kernel_size=2, stride=1, padding=0)\n",
       "        (relu): ReLU(inplace)\n",
       "      )\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (relu): ReLU()\n",
       "    )\n",
       "    (6): UnetBlock(\n",
       "      (shuf): PixelShuffle_ICNR(\n",
       "        (conv): Sequential(\n",
       "          (0): Conv2d(384, 768, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (shuf): PixelShuffle(upscale_factor=2)\n",
       "        (pad): ReplicationPad2d((1, 0, 1, 0))\n",
       "        (blur): AvgPool2d(kernel_size=2, stride=1, padding=0)\n",
       "        (relu): ReLU(inplace)\n",
       "      )\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (relu): ReLU()\n",
       "    )\n",
       "    (7): UnetBlock(\n",
       "      (shuf): PixelShuffle_ICNR(\n",
       "        (conv): Sequential(\n",
       "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (shuf): PixelShuffle(upscale_factor=2)\n",
       "        (pad): ReplicationPad2d((1, 0, 1, 0))\n",
       "        (blur): AvgPool2d(kernel_size=2, stride=1, padding=0)\n",
       "        (relu): ReLU(inplace)\n",
       "      )\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (relu): ReLU()\n",
       "    )\n",
       "    (8): UnetBlock(\n",
       "      (shuf): PixelShuffle_ICNR(\n",
       "        (conv): Sequential(\n",
       "          (0): Conv2d(192, 384, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (shuf): PixelShuffle(upscale_factor=2)\n",
       "        (pad): ReplicationPad2d((1, 0, 1, 0))\n",
       "        (blur): AvgPool2d(kernel_size=2, stride=1, padding=0)\n",
       "        (relu): ReLU(inplace)\n",
       "      )\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(160, 80, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(80, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (relu): ReLU()\n",
       "    )\n",
       "    (9): PixelShuffle_ICNR(\n",
       "      (conv): Sequential(\n",
       "        (0): Conv2d(80, 320, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (shuf): PixelShuffle(upscale_factor=2)\n",
       "      (pad): ReplicationPad2d((1, 0, 1, 0))\n",
       "      (blur): AvgPool2d(kernel_size=2, stride=1, padding=0)\n",
       "      (relu): ReLU(inplace)\n",
       "    )\n",
       "    (10): MergeLayer()\n",
       "    (11): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(83, 83, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(83, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(83, 83, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(83, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): MergeLayer()\n",
       "      )\n",
       "    )\n",
       "    (12): Sequential(\n",
       "      (0): Conv2d(83, 10, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (1): BatchNorm2d(10, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchsummary import summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 112, 112]           9,408\n",
      "              ReLU-2         [-1, 64, 112, 112]               0\n",
      "       BatchNorm2d-3         [-1, 64, 112, 112]             128\n",
      "         MaxPool2d-4           [-1, 64, 56, 56]               0\n",
      "            Conv2d-5           [-1, 64, 56, 56]          36,864\n",
      "              ReLU-6           [-1, 64, 56, 56]               0\n",
      "       BatchNorm2d-7           [-1, 64, 56, 56]             128\n",
      "            Conv2d-8           [-1, 64, 56, 56]          36,864\n",
      "              ReLU-9           [-1, 64, 56, 56]               0\n",
      "      BatchNorm2d-10           [-1, 64, 56, 56]             128\n",
      "       MergeLayer-11           [-1, 64, 56, 56]               0\n",
      "     SequentialEx-12           [-1, 64, 56, 56]               0\n",
      "           Conv2d-13          [-1, 128, 28, 28]          73,728\n",
      "             ReLU-14          [-1, 128, 28, 28]               0\n",
      "      BatchNorm2d-15          [-1, 128, 28, 28]             256\n",
      "           Conv2d-16          [-1, 128, 28, 28]         147,456\n",
      "             ReLU-17          [-1, 128, 28, 28]               0\n",
      "      BatchNorm2d-18          [-1, 128, 28, 28]             256\n",
      "       MergeLayer-19          [-1, 128, 28, 28]               0\n",
      "     SequentialEx-20          [-1, 128, 28, 28]               0\n",
      "     SequentialEx-21          [-1, 128, 28, 28]               0\n",
      "           Conv2d-22          [-1, 256, 14, 14]         294,912\n",
      "             ReLU-23          [-1, 256, 14, 14]               0\n",
      "      BatchNorm2d-24          [-1, 256, 14, 14]             512\n",
      "           Conv2d-25          [-1, 256, 14, 14]         589,824\n",
      "             ReLU-26          [-1, 256, 14, 14]               0\n",
      "      BatchNorm2d-27          [-1, 256, 14, 14]             512\n",
      "       MergeLayer-28          [-1, 256, 14, 14]               0\n",
      "     SequentialEx-29          [-1, 256, 14, 14]               0\n",
      "     SequentialEx-30          [-1, 256, 14, 14]               0\n",
      "           Conv2d-31            [-1, 512, 7, 7]       1,179,648\n",
      "             ReLU-32            [-1, 512, 7, 7]               0\n",
      "      BatchNorm2d-33            [-1, 512, 7, 7]           1,024\n",
      "           Conv2d-34            [-1, 512, 7, 7]       2,359,296\n",
      "             ReLU-35            [-1, 512, 7, 7]               0\n",
      "      BatchNorm2d-36            [-1, 512, 7, 7]           1,024\n",
      "       MergeLayer-37            [-1, 512, 7, 7]               0\n",
      "     SequentialEx-38            [-1, 512, 7, 7]               0\n",
      "     SequentialEx-39            [-1, 512, 7, 7]               0\n",
      "     SequentialEx-40            [-1, 512, 7, 7]               0\n",
      "      BatchNorm2d-41            [-1, 512, 7, 7]           1,024\n",
      "             ReLU-42            [-1, 512, 7, 7]               0\n",
      "           Conv2d-43           [-1, 1024, 7, 7]       4,718,592\n",
      "             ReLU-44           [-1, 1024, 7, 7]               0\n",
      "      BatchNorm2d-45           [-1, 1024, 7, 7]           2,048\n",
      "           Conv2d-46            [-1, 512, 7, 7]       4,718,592\n",
      "             ReLU-47            [-1, 512, 7, 7]               0\n",
      "      BatchNorm2d-48            [-1, 512, 7, 7]           1,024\n",
      "           Conv2d-49           [-1, 1024, 7, 7]         525,312\n",
      "             ReLU-50           [-1, 1024, 7, 7]               0\n",
      "     PixelShuffle-51          [-1, 256, 14, 14]               0\n",
      "PixelShuffle_ICNR-52          [-1, 256, 14, 14]               0\n",
      "      BatchNorm2d-53          [-1, 256, 14, 14]             512\n",
      "             ReLU-54          [-1, 512, 14, 14]               0\n",
      "           Conv2d-55          [-1, 512, 14, 14]       2,359,296\n",
      "             ReLU-56          [-1, 512, 14, 14]               0\n",
      "      BatchNorm2d-57          [-1, 512, 14, 14]           1,024\n",
      "        UnetBlock-58          [-1, 512, 14, 14]               0\n",
      "           Conv2d-59         [-1, 1024, 14, 14]         525,312\n",
      "             ReLU-60         [-1, 1024, 14, 14]               0\n",
      "     PixelShuffle-61          [-1, 256, 28, 28]               0\n",
      "PixelShuffle_ICNR-62          [-1, 256, 28, 28]               0\n",
      "      BatchNorm2d-63          [-1, 128, 28, 28]             256\n",
      "             ReLU-64          [-1, 384, 28, 28]               0\n",
      "           Conv2d-65          [-1, 384, 28, 28]       1,327,104\n",
      "             ReLU-66          [-1, 384, 28, 28]               0\n",
      "      BatchNorm2d-67          [-1, 384, 28, 28]             768\n",
      "        UnetBlock-68          [-1, 384, 28, 28]               0\n",
      "           Conv2d-69          [-1, 768, 28, 28]         295,680\n",
      "             ReLU-70          [-1, 768, 28, 28]               0\n",
      "     PixelShuffle-71          [-1, 192, 56, 56]               0\n",
      "PixelShuffle_ICNR-72          [-1, 192, 56, 56]               0\n",
      "      BatchNorm2d-73           [-1, 64, 56, 56]             128\n",
      "             ReLU-74          [-1, 256, 56, 56]               0\n",
      "           Conv2d-75          [-1, 256, 56, 56]         589,824\n",
      "             ReLU-76          [-1, 256, 56, 56]               0\n",
      "      BatchNorm2d-77          [-1, 256, 56, 56]             512\n",
      "        UnetBlock-78          [-1, 256, 56, 56]               0\n",
      "           Conv2d-79          [-1, 512, 56, 56]         131,584\n",
      "             ReLU-80          [-1, 512, 56, 56]               0\n",
      "     PixelShuffle-81        [-1, 128, 112, 112]               0\n",
      "PixelShuffle_ICNR-82        [-1, 128, 112, 112]               0\n",
      "      BatchNorm2d-83         [-1, 64, 112, 112]             128\n",
      "             ReLU-84        [-1, 192, 112, 112]               0\n",
      "           Conv2d-85        [-1, 192, 112, 112]         331,776\n",
      "             ReLU-86        [-1, 192, 112, 112]               0\n",
      "      BatchNorm2d-87        [-1, 192, 112, 112]             384\n",
      "        UnetBlock-88        [-1, 192, 112, 112]               0\n",
      "           Conv2d-89        [-1, 384, 112, 112]          74,112\n",
      "             ReLU-90        [-1, 384, 112, 112]               0\n",
      "     PixelShuffle-91         [-1, 96, 224, 224]               0\n",
      "PixelShuffle_ICNR-92         [-1, 96, 224, 224]               0\n",
      "      BatchNorm2d-93         [-1, 64, 112, 112]             128\n",
      "             ReLU-94        [-1, 160, 112, 112]               0\n",
      "           Conv2d-95         [-1, 80, 112, 112]         115,200\n",
      "             ReLU-96         [-1, 80, 112, 112]               0\n",
      "      BatchNorm2d-97         [-1, 80, 112, 112]             160\n",
      "        UnetBlock-98         [-1, 80, 112, 112]               0\n",
      "           Conv2d-99        [-1, 320, 112, 112]          25,920\n",
      "            ReLU-100        [-1, 320, 112, 112]               0\n",
      "    PixelShuffle-101         [-1, 80, 224, 224]               0\n",
      "PixelShuffle_ICNR-102         [-1, 80, 224, 224]               0\n",
      "      MergeLayer-103         [-1, 83, 224, 224]               0\n",
      "          Conv2d-104         [-1, 83, 224, 224]          62,001\n",
      "            ReLU-105         [-1, 83, 224, 224]               0\n",
      "     BatchNorm2d-106         [-1, 83, 224, 224]             166\n",
      "          Conv2d-107         [-1, 83, 224, 224]          62,001\n",
      "            ReLU-108         [-1, 83, 224, 224]               0\n",
      "     BatchNorm2d-109         [-1, 83, 224, 224]             166\n",
      "      MergeLayer-110         [-1, 83, 224, 224]               0\n",
      "    SequentialEx-111         [-1, 83, 224, 224]               0\n",
      "          Conv2d-112         [-1, 10, 224, 224]             830\n",
      "     BatchNorm2d-113         [-1, 10, 224, 224]              20\n",
      "================================================================\n",
      "Total params: 20,603,552\n",
      "Trainable params: 20,603,552\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.57\n",
      "Forward/backward pass size (MB): 883.34\n",
      "Params size (MB): 78.60\n",
      "Estimated Total Size (MB): 962.51\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "summary(unet, input_size=(3, 224, 224))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.vision.models import DynamicUnet as DUnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "new = DUnet(create_body(models.resnet34), n_classes=10).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.vision import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DynamicUnet(\n",
       "  (layers): ModuleList(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): ReLU(inplace)\n",
       "      (3): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "      (4): Sequential(\n",
       "        (0): BasicBlock(\n",
       "          (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): BasicBlock(\n",
       "          (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): BasicBlock(\n",
       "          (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (5): Sequential(\n",
       "        (0): BasicBlock(\n",
       "          (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (1): BasicBlock(\n",
       "          (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): BasicBlock(\n",
       "          (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (3): BasicBlock(\n",
       "          (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (6): Sequential(\n",
       "        (0): BasicBlock(\n",
       "          (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (1): BasicBlock(\n",
       "          (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): BasicBlock(\n",
       "          (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (3): BasicBlock(\n",
       "          (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (4): BasicBlock(\n",
       "          (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (5): BasicBlock(\n",
       "          (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (7): Sequential(\n",
       "        (0): BasicBlock(\n",
       "          (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "            (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (1): BasicBlock(\n",
       "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): BasicBlock(\n",
       "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU(inplace)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU()\n",
       "    (3): Sequential(\n",
       "      (0): Sequential(\n",
       "        (0): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (1): Sequential(\n",
       "        (0): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (4): UnetBlock(\n",
       "      (shuf): PixelShuffle_ICNR(\n",
       "        (conv): Sequential(\n",
       "          (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (shuf): PixelShuffle(upscale_factor=2)\n",
       "        (pad): ReplicationPad2d((1, 0, 1, 0))\n",
       "        (blur): AvgPool2d(kernel_size=2, stride=1, padding=0)\n",
       "        (relu): ReLU(inplace)\n",
       "      )\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (relu): ReLU()\n",
       "    )\n",
       "    (5): UnetBlock(\n",
       "      (shuf): PixelShuffle_ICNR(\n",
       "        (conv): Sequential(\n",
       "          (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (shuf): PixelShuffle(upscale_factor=2)\n",
       "        (pad): ReplicationPad2d((1, 0, 1, 0))\n",
       "        (blur): AvgPool2d(kernel_size=2, stride=1, padding=0)\n",
       "        (relu): ReLU(inplace)\n",
       "      )\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (relu): ReLU()\n",
       "    )\n",
       "    (6): UnetBlock(\n",
       "      (shuf): PixelShuffle_ICNR(\n",
       "        (conv): Sequential(\n",
       "          (0): Conv2d(384, 768, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (shuf): PixelShuffle(upscale_factor=2)\n",
       "        (pad): ReplicationPad2d((1, 0, 1, 0))\n",
       "        (blur): AvgPool2d(kernel_size=2, stride=1, padding=0)\n",
       "        (relu): ReLU(inplace)\n",
       "      )\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (relu): ReLU()\n",
       "    )\n",
       "    (7): UnetBlock(\n",
       "      (shuf): PixelShuffle_ICNR(\n",
       "        (conv): Sequential(\n",
       "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (shuf): PixelShuffle(upscale_factor=2)\n",
       "        (pad): ReplicationPad2d((1, 0, 1, 0))\n",
       "        (blur): AvgPool2d(kernel_size=2, stride=1, padding=0)\n",
       "        (relu): ReLU(inplace)\n",
       "      )\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv1): Sequential(\n",
       "        (0): Conv2d(192, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (conv2): Sequential(\n",
       "        (0): Conv2d(96, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): ReLU(inplace)\n",
       "        (2): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (relu): ReLU()\n",
       "    )\n",
       "    (8): PixelShuffle_ICNR(\n",
       "      (conv): Sequential(\n",
       "        (0): Conv2d(96, 384, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (shuf): PixelShuffle(upscale_factor=2)\n",
       "      (pad): ReplicationPad2d((1, 0, 1, 0))\n",
       "      (blur): AvgPool2d(kernel_size=2, stride=1, padding=0)\n",
       "      (relu): ReLU(inplace)\n",
       "    )\n",
       "    (9): MergeLayer()\n",
       "    (10): SequentialEx(\n",
       "      (layers): ModuleList(\n",
       "        (0): Sequential(\n",
       "          (0): Conv2d(99, 99, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(99, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): Sequential(\n",
       "          (0): Conv2d(99, 99, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): ReLU(inplace)\n",
       "          (2): BatchNorm2d(99, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (2): MergeLayer()\n",
       "      )\n",
       "    )\n",
       "    (11): Sequential(\n",
       "      (0): Conv2d(99, 10, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (1): BatchNorm2d(10, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchviz import make_dot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<generator object Module.parameters at 0x7fc3d5f6b6d0>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new.parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = net(torch.randn(1, 3, 224, 224).cuda())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "ename": "ExecutableNotFound",
     "evalue": "failed to execute ['dot', '-Tsvg'], make sure the Graphviz executables are on your systems' PATH",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m~/.conda/envs/fastai/lib/python3.6/site-packages/graphviz/backend.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(cmd, input, capture_output, check, quiet, **kwargs)\u001b[0m\n\u001b[1;32m    157\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 158\u001b[0;31m         \u001b[0mproc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msubprocess\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mPopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcmd\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstartupinfo\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mget_startupinfo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    159\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mOSError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/fastai/lib/python3.6/subprocess.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, args, bufsize, executable, stdin, stdout, stderr, preexec_fn, close_fds, shell, cwd, env, universal_newlines, startupinfo, creationflags, restore_signals, start_new_session, pass_fds, encoding, errors)\u001b[0m\n\u001b[1;32m    728\u001b[0m                                 \u001b[0merrread\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0merrwrite\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 729\u001b[0;31m                                 restore_signals, start_new_session)\n\u001b[0m\u001b[1;32m    730\u001b[0m         \u001b[0;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/fastai/lib/python3.6/subprocess.py\u001b[0m in \u001b[0;36m_execute_child\u001b[0;34m(self, args, executable, preexec_fn, close_fds, pass_fds, cwd, env, startupinfo, creationflags, shell, p2cread, p2cwrite, c2pread, c2pwrite, errread, errwrite, restore_signals, start_new_session)\u001b[0m\n\u001b[1;32m   1363\u001b[0m                             \u001b[0merr_msg\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;34m': '\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mrepr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0merr_filename\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1364\u001b[0;31m                     \u001b[0;32mraise\u001b[0m \u001b[0mchild_exception_type\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0merrno_num\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0merr_msg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0merr_filename\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1365\u001b[0m                 \u001b[0;32mraise\u001b[0m \u001b[0mchild_exception_type\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0merr_msg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'dot': 'dot'",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mExecutableNotFound\u001b[0m                        Traceback (most recent call last)",
      "\u001b[0;32m~/.conda/envs/fastai/lib/python3.6/site-packages/IPython/core/formatters.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, obj)\u001b[0m\n\u001b[1;32m    343\u001b[0m             \u001b[0mmethod\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_real_method\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprint_method\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    344\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mmethod\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 345\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mmethod\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    346\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    347\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/fastai/lib/python3.6/site-packages/graphviz/files.py\u001b[0m in \u001b[0;36m_repr_svg_\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    111\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    112\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_repr_svg_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 113\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpipe\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'svg'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_encoding\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    114\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    115\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mpipe\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mformat\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrenderer\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mformatter\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mquiet\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/fastai/lib/python3.6/site-packages/graphviz/files.py\u001b[0m in \u001b[0;36mpipe\u001b[0;34m(self, format, renderer, formatter, quiet)\u001b[0m\n\u001b[1;32m    136\u001b[0m         out = backend.pipe(self._engine, format, data,\n\u001b[1;32m    137\u001b[0m                            \u001b[0mrenderer\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mrenderer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mformatter\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mformatter\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 138\u001b[0;31m                            quiet=quiet)\n\u001b[0m\u001b[1;32m    139\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    140\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/fastai/lib/python3.6/site-packages/graphviz/backend.py\u001b[0m in \u001b[0;36mpipe\u001b[0;34m(engine, format, data, renderer, formatter, quiet)\u001b[0m\n\u001b[1;32m    226\u001b[0m     \"\"\"\n\u001b[1;32m    227\u001b[0m     \u001b[0mcmd\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcommand\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mengine\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mformat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrenderer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mformatter\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 228\u001b[0;31m     \u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcmd\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcapture_output\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcheck\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mquiet\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mquiet\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    229\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    230\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/fastai/lib/python3.6/site-packages/graphviz/backend.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(cmd, input, capture_output, check, quiet, **kwargs)\u001b[0m\n\u001b[1;32m    159\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mOSError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    160\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0merrno\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0merrno\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mENOENT\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 161\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0mExecutableNotFound\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcmd\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    162\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    163\u001b[0m             \u001b[0;32mraise\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mExecutableNotFound\u001b[0m: failed to execute ['dot', '-Tsvg'], make sure the Graphviz executables are on your systems' PATH"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<graphviz.dot.Digraph at 0x7fc334d881d0>"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "make_dot(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models.custom_resnet import _resnet, BasicBlock\n",
    "\n",
    "model = _resnet('resnet14', BasicBlock, [2, 2, 1, 1], pretrained = False, progress = False).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 112, 112]           9,408\n",
      "       BatchNorm2d-2         [-1, 64, 112, 112]             128\n",
      "              ReLU-3         [-1, 64, 112, 112]               0\n",
      "         MaxPool2d-4           [-1, 64, 56, 56]               0\n",
      "            Conv2d-5           [-1, 64, 56, 56]          36,864\n",
      "       BatchNorm2d-6           [-1, 64, 56, 56]             128\n",
      "              ReLU-7           [-1, 64, 56, 56]               0\n",
      "            Conv2d-8           [-1, 64, 56, 56]          36,864\n",
      "       BatchNorm2d-9           [-1, 64, 56, 56]             128\n",
      "             ReLU-10           [-1, 64, 56, 56]               0\n",
      "       BasicBlock-11           [-1, 64, 56, 56]               0\n",
      "           Conv2d-12           [-1, 64, 56, 56]          36,864\n",
      "      BatchNorm2d-13           [-1, 64, 56, 56]             128\n",
      "             ReLU-14           [-1, 64, 56, 56]               0\n",
      "           Conv2d-15           [-1, 64, 56, 56]          36,864\n",
      "      BatchNorm2d-16           [-1, 64, 56, 56]             128\n",
      "             ReLU-17           [-1, 64, 56, 56]               0\n",
      "       BasicBlock-18           [-1, 64, 56, 56]               0\n",
      "           Conv2d-19          [-1, 128, 28, 28]          73,728\n",
      "      BatchNorm2d-20          [-1, 128, 28, 28]             256\n",
      "             ReLU-21          [-1, 128, 28, 28]               0\n",
      "           Conv2d-22          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-23          [-1, 128, 28, 28]             256\n",
      "           Conv2d-24          [-1, 128, 28, 28]           8,192\n",
      "      BatchNorm2d-25          [-1, 128, 28, 28]             256\n",
      "             ReLU-26          [-1, 128, 28, 28]               0\n",
      "       BasicBlock-27          [-1, 128, 28, 28]               0\n",
      "           Conv2d-28          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-29          [-1, 128, 28, 28]             256\n",
      "             ReLU-30          [-1, 128, 28, 28]               0\n",
      "           Conv2d-31          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-32          [-1, 128, 28, 28]             256\n",
      "             ReLU-33          [-1, 128, 28, 28]               0\n",
      "       BasicBlock-34          [-1, 128, 28, 28]               0\n",
      "           Conv2d-35          [-1, 256, 14, 14]         294,912\n",
      "      BatchNorm2d-36          [-1, 256, 14, 14]             512\n",
      "             ReLU-37          [-1, 256, 14, 14]               0\n",
      "           Conv2d-38          [-1, 256, 14, 14]         589,824\n",
      "      BatchNorm2d-39          [-1, 256, 14, 14]             512\n",
      "           Conv2d-40          [-1, 256, 14, 14]          32,768\n",
      "      BatchNorm2d-41          [-1, 256, 14, 14]             512\n",
      "             ReLU-42          [-1, 256, 14, 14]               0\n",
      "       BasicBlock-43          [-1, 256, 14, 14]               0\n",
      "           Conv2d-44            [-1, 512, 7, 7]       1,179,648\n",
      "      BatchNorm2d-45            [-1, 512, 7, 7]           1,024\n",
      "             ReLU-46            [-1, 512, 7, 7]               0\n",
      "           Conv2d-47            [-1, 512, 7, 7]       2,359,296\n",
      "      BatchNorm2d-48            [-1, 512, 7, 7]           1,024\n",
      "           Conv2d-49            [-1, 512, 7, 7]         131,072\n",
      "      BatchNorm2d-50            [-1, 512, 7, 7]           1,024\n",
      "             ReLU-51            [-1, 512, 7, 7]               0\n",
      "       BasicBlock-52            [-1, 512, 7, 7]               0\n",
      "AdaptiveMaxPool2d-53            [-1, 512, 1, 1]               0\n",
      "AdaptiveAvgPool2d-54            [-1, 512, 1, 1]               0\n",
      "AdaptiveConcatPool2d-55           [-1, 1024, 1, 1]               0\n",
      "           Linear-56                  [-1, 256]         262,400\n",
      "             ReLU-57                  [-1, 256]               0\n",
      "           Linear-58                   [-1, 10]           2,570\n",
      "================================================================\n",
      "Total params: 5,540,170\n",
      "Trainable params: 5,540,170\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.57\n",
      "Forward/backward pass size (MB): 58.78\n",
      "Params size (MB): 21.13\n",
      "Estimated Total Size (MB): 80.49\n",
      "----------------------------------------------------------------\n",
      "None\n",
      "ResNet(\n",
      "  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
      "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (relu): ReLU(inplace)\n",
      "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
      "  (layer1): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    )\n",
      "  )\n",
      "  (layer2): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    )\n",
      "  )\n",
      "  (layer3): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (layer4): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (avgpool): AdaptiveConcatPool2d(\n",
      "    (ap): AdaptiveAvgPool2d(output_size=1)\n",
      "    (mp): AdaptiveMaxPool2d(output_size=1)\n",
      "  )\n",
      "  (fc): Linear(in_features=1024, out_features=256, bias=True)\n",
      "  (relu2): ReLU(inplace)\n",
      "  (fc2): Linear(in_features=256, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from torchsummary import summary\n",
    "print(summary(model, (3, 224, 224)))\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (ak_fastai)",
   "language": "python",
   "name": "ak_fastai"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
